import torch
import torch.nn as nn
from torchvision import transforms
import cv2
import numpy as np
from PIL import Image

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


class lenet(nn.Module):
    def __init__(self):
        super().__init__()
        self.dnn = torch.nn.Sequential(
            # Net_1
            #nn.Conv2d(1, 6, 5, padding=2), nn.Sigmoid(),
            #nn.MaxPool2d(2, 2),
            #nn.Conv2d(6, 16, 5), nn.Sigmoid(),
            #nn.MaxPool2d(2, 2), nn.Flatten(),
            #nn.Linear(16 * 5 * 5, 120), nn.Sigmoid(),
            #nn.Linear(120, 84), nn.Sigmoid(), nn.Linear(84, 10)

            # Net_2
            # nn.Conv2d(1, 6, 5, padding=2), nn.Sigmoid(),
            # nn.AvgPool2d(2, 2),
            # nn.Conv2d(6, 16, 5), nn.Sigmoid(),
            # nn.AvgPool2d(2, 2), nn.Flatten(),
            # nn.Linear(16 * 5 * 5, 120), nn.Sigmoid(),
            # nn.Linear(120, 84), nn.Sigmoid(), nn.Linear(84, 10)

            # Net_3
            nn.Conv2d(1, 6, 5, padding=2), nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5), nn.ReLU(),
            nn.MaxPool2d(2, 2), nn.Flatten(),
            nn.Linear(16 * 5 * 5, 120), nn.ReLU(),
            nn.Linear(120, 84), nn.ReLU(),
            nn.Linear(84, 10)

            # Net_5
            # nn.Conv2d(1, 6, 5, padding=2), nn.ReLU(),
            # nn.MaxPool2d(2, 2),
            # nn.Conv2d(6, 16, 5), nn.ReLU(),
            # nn.MaxPool2d(2, 2), nn.Flatten(),
            # nn.Linear(16 * 5 * 5, 120), nn.ReLU(),
            # nn.Linear(120, 84), nn.ReLU(),
            # nn.Linear(84, 50), nn.ReLU(),
            # nn.Linear(50, 10)
            )

    def forward(self, x):
        out = self.dnn(x)
        return out


def process(img, i, j, count):
    # 出口
    if i < 0 or j < 0 or i > img.shape[0]-1 or j > img.shape[1]-1 or img[i][j] != 255:
        return
    # 感染成count值
    img[i][j] = count
    # 递归
    process(img, i, j+1, count)
    process(img, i, j-1, count)
    process(img, i+1, j, count)
    process(img, i-1, j, count)


def numisland(img):
    count = 1  # 数字的像素值修改
    row = img.shape[0]
    col = img.shape[1]
    for j in range(col):
        for i in range(row):
            if img[i][j] == 255:  # 遇到了数字
                # 进入感染
                process(img, i, j, count)
                count += 1
    return img, count


def foundnum(img_binary, count):
    # 截取数字
    x_list = []
    y_list = []
    for i in range(0, img_binary.shape[0]):
        for j in range(0, img_binary.shape[1]):
            if img_binary[i, j] == count:
                img_binary[i, j] = 255
                x_list.append(i)
                y_list.append(j)

    # 选出含有数字的区域
    x_min = min(x_list)
    x_max = max(x_list)
    y_min = min(y_list)
    y_max = max(y_list)
    x_center = int((x_min + x_max) / 2)
    y_center = int((y_min + y_max) / 2)
    longer_side = max([x_max - x_min, y_max - y_min])
    # 边缘多出一点
    s_d = int(longer_side / 5)
    img_binary = img_binary[x_center - int(longer_side/2) - s_d:x_center + int(longer_side/2) + s_d, y_center - int(longer_side/2) - s_d:y_center + int(longer_side/2) + s_d]
    print("二值剪切之后形状：", img_binary.shape)

    img_binay_resize = cv2.resize(img_binary, (28, 28), interpolation=cv2.INTER_AREA)
    print("二值剪切变小之后形状：", img_binay_resize.shape)

    # 转成tensor类型
    transf = transforms.ToTensor()
    my_mnist = transf(img_binay_resize)
    # 变成tensor类型的四维张量
    my_mnist = my_mnist.view(1, 1, 28, 28)

    # 送入GPU
    input = my_mnist.to(device)
    output = net(input)
    # print('0~9的预测值:', output)  # 得到0~9的预测值

    # 把tensor类型转换为numpy类型
    output = output.detach().cpu().numpy()
    # np.argmax(a) 取出a中元素最大值所对应的索引，即取出预测值最高的数字
    pred = np.argmax(output[0])
    # print('预测数字：', pred)     # , '   预测值：', output[0][pred]
    return pred


net = lenet()
# 加载训练好的模型
net.load_state_dict(torch.load('model/net_3.pth'))
# 送进GPU
net = net.to(device)
# 读取要预测的图片
img = cv2.imread('test1.png', 0)
down_width = 300
down_height = 100
down_points = (down_width, down_height)
img = cv2.resize(img, down_points, interpolation=cv2.INTER_LINEAR)

# 阈值函数————将图像的像素点经过阈值（threshold）比较重新设置为0或者255,大于阈值的部分被置为255,小于部分被置为0
ret, img_binary = cv2.threshold(img, 90, 255, cv2.THRESH_BINARY)
# 反转成黑底白字
img_binary = 255 - img_binary
# cv2.imwrite('二值图像.png', img_binary)
# cv2.imshow('img_binary', img_binary)
print(img_binary.shape)


# 数字列表，输出检测到的数字
numlist = []

# 感染算法，把数字区分开 （底0，第一个数字的255改为1，第二个是2，以此类推）
img_island, count = numisland(img_binary)
for x in range(1, count):
    pred = foundnum(img_island, x)
    numlist.append(pred)

# 输出预测结果
print(numlist)
